"""
主程序入口
"""
import os
import sys
from datetime import datetime, timedelta
from typing import Dict, List

# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.utils.config import Config
from src.utils.logger import get_logger
from src.okx.client import OKXClient
from src.data.kline import KlineDataFetcher
from src.data.processor import KlineProcessor
from src.strategy.scalping import ScalpingStrategy
from src.trading.executor import TradeExecutor
from src.trading.risk import RiskManager, RiskControl
from src.backtesting.engine import BacktestEngine
from src.simulator.mock_exchange import MockExchange
from src.costs.calculator import CostCalculator

logger = get_logger(__name__)


class ScalpingBot:
    """头皮策略交易机器人"""

    def __init__(self, config_path: str = None):
        """初始化机器人"""
        # 加载配置
        self.config = Config.load(config_path)

        # 初始化API客户端
        self.client = OKXClient(
            api_key=self.config.get('api.api_key', ''),
            secret_key=self.config.get('api.secret_key', ''),
            passphrase=self.config.get('api.passphrase', ''),
            environment=self.config.get('api.environment', 'sandbox'),
            simulate=self.config.get('simulate.enabled', False)
        )

        # 初始化数据处理
        self.fetcher = KlineDataFetcher(self.client)
        self.processor = KlineProcessor()

        # 初始化策略
        self.strategy = ScalpingStrategy(self.config.get('strategy', {}))

        # 初始化交易执行
        self.executor = TradeExecutor(
            self.client,
            trade_type=self.config.get('trading.trade_type', 'spot'),
            leverage=self.config.get('trading.leverage', 1.0),
            cost_calculator=CostCalculator(
                maker_fee=self.config.get('costs.maker_fee', 0.02),
                taker_fee=self.config.get('costs.taker_fee', 0.05),
                funding_rate=self.config.get('costs.funding_rate', 0.03)
            )
        )

        # 初始化风险管理
        risk_control = RiskControl(
            stop_loss_pct=self.config.get('risk.stop_loss', 1.0),
            take_profit_pct=self.config.get('risk.take_profit', 2.0),
            max_positions=self.config.get('risk.max_positions', 3),
            max_daily_loss_pct=self.config.get('risk.max_daily_loss', 5.0),
            max_drawdown_pct=self.config.get('risk.max_drawdown', 10.0),
            max_position_bars=self.config.get('trading.max_position_bars', 20),
        )
        self.risk_manager = RiskManager(risk_control, self.config.get('trading.position_size', 100))

        logger.info("机器人初始化完成")

    def run_backtest(self):
        """运行回测"""
        logger.info("="*60)
        logger.info("启动回测模式")
        logger.info("="*60)

        # 获取回测配置
        backtest_config = self.config.get('backtest', {})
        start_date = datetime.strptime(backtest_config['start_date'], '%Y-%m-%d')
        end_date = datetime.strptime(backtest_config['end_date'], '%Y-%m-%d')
        symbols = backtest_config.get('symbols', ['BTC-USDT'])
        initial_balance = backtest_config.get('initial_balance', 10000)

        # 初始化回测引擎
        backtest_engine = BacktestEngine(self.config, initial_balance)

        # 获取历史K线数据
        klines_data = {}
        for symbol in symbols:
            logger.info(f"获取历史K线: {symbol}")

            klines_data[symbol] = {}
            for timeframe in self.config.get('strategy.timeframes', [5, 15, 60, 240]):
                klines = self.fetcher.fetch_historical_klines(symbol, timeframe, start_date, end_date)
                klines_data[symbol][timeframe] = klines

        # 运行回测
        report = backtest_engine.run_backtest(symbols, klines_data, start_date, end_date)

        # 保存报告
        self._save_backtest_report(report)

        return report

    def run_live(self):
        """运行实盘交易"""
        logger.info("="*60)
        logger.info("启动实盘交易模式")
        logger.info("="*60)

        symbols = self.config.get('trading.symbols', ['BTC-USDT', 'ETH-USDT'])
        timeframes = self.config.get('strategy.timeframes', [5, 15, 60, 240])

        try:
            while True:
                for symbol in symbols:
                    # 获取最新K线
                    kline_data = {}
                    for tf in timeframes:
                        klines = self.fetcher.fetch_klines(symbol, tf, limit=100)
                        if klines:
                            self.processor.add_klines(symbol, tf, klines)
                            df = self.processor.get_klines(symbol, tf)
                            df = self.processor.add_technical_indicators(df, self.config.get('strategy', {}))
                            kline_data[tf] = df

                    # 生成信号
                    signal, confidence = self.strategy.generate_signal(kline_data)

                    if signal in ["buy", "sell"] and confidence > 0.6:
                        logger.info(f"信号: {signal.upper()}, {symbol}, 信心度: {confidence:.2f}")

                        # 检查现仓位
                        position = self.executor.position_manager.get_active_position(symbol)

                        if position and signal != "hold":
                            # 平仓
                            self.executor.close_position(symbol, reason="signal")
                        elif not position and signal in ["buy", "sell"]:
                            # 开仓
                            position_size = self.config.get('trading.position_size', 100)
                            quantity = position_size / self.processor.get_klines(symbol, 5)['close'].iloc[-1]

                            if signal == "buy":
                                self.executor.execute_buy(symbol, quantity, order_type="market")
                            else:
                                self.executor.execute_short(symbol, quantity, order_type="market")

                    # 检查活跃持仓
                    for trade in self.executor.position_manager.get_active_positions():
                        # 检查K线限制
                        bars_held = (datetime.now() - trade.entry_time).total_seconds() / (5 * 60)
                        if bars_held >= self.config.get('trading.max_position_bars', 20):
                            self.executor.close_position(trade.symbol, reason="timeout")

                # 每分钟检查一次
                logger.debug("等待下一个检查周期...")
                import time
                time.sleep(60)

        except KeyboardInterrupt:
            logger.info("收到中断信号，关闭机器人")
        except Exception as e:
            logger.error(f"交易异常: {str(e)}")

    def run_simulate(self):
        """运行模拟盘交易"""
        logger.info("="*60)
        logger.info("启动模拟盘交易模式")
        logger.info("="*60)

        # 初始化模拟交易所
        initial_balances = {
            'USDT': self.config.get('simulate.initial_balance', 5000),
            'BTC': 0,
            'ETH': 0,
        }
        mock_exchange = MockExchange(initial_balances)

        symbols = self.config.get('trading.symbols', ['BTC-USDT', 'ETH-USDT'])
        timeframes = self.config.get('strategy.timeframes', [5, 15, 60, 240])

        logger.info(f"模拟初始余额: {initial_balances}")

        try:
            iteration = 0
            while True:
                iteration += 1
                logger.info(f"--- 迭代 {iteration} ---")

                for symbol in symbols:
                    # 获取最新K线
                    kline_data = {}
                    for tf in timeframes:
                        klines = self.fetcher.fetch_klines(symbol, tf, limit=100)
                        if klines:
                            self.processor.add_klines(symbol, tf, klines)
                            df = self.processor.get_klines(symbol, tf)
                            df = self.processor.add_technical_indicators(df, self.config.get('strategy', {}))
                            kline_data[tf] = df

                    # 设置当前价格
                    if 5 in kline_data and len(kline_data[5]) > 0:
                        current_price = kline_data[5]['close'].iloc[-1]
                        mock_exchange.set_ticker(symbol, current_price)

                    # 生成信号
                    signal, confidence = self.strategy.generate_signal(kline_data)

                    if signal in ["buy", "sell"] and confidence > 0.6:
                        logger.info(f"信号: {signal.upper()}, {symbol}, 信心度: {confidence:.2f}")

                        position_size = self.config.get('trading.position_size', 100)
                        quantity = position_size / current_price if current_price > 0 else 0

                        if signal == "buy":
                            result = mock_exchange.place_order(symbol, "buy", "market", quantity)
                            if result.success:
                                logger.info(f"买入成功: {symbol}, 数量={quantity}")

                        elif signal == "sell":
                            result = mock_exchange.place_order(symbol, "sell", "market", quantity)
                            if result.success:
                                logger.info(f"卖出成功: {symbol}, 数量={quantity}")

                # 显示账户余额
                balances = mock_exchange.get_balance()
                total_value = sum(b.balance for b in balances)
                logger.info(f"账户价值: {total_value:.2f} USDT")

                # 等待下一个周期
                import time
                time.sleep(10)

        except KeyboardInterrupt:
            logger.info("收到中断信号，关闭模拟盘")

    def _save_backtest_report(self, report: Dict):
        """保存回测报告"""
        import json

        os.makedirs('reports', exist_ok=True)
        filename = f"reports/backtest_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"

        with open(filename, 'w') as f:
            json.dump(report, f, indent=2, default=str)

        logger.info(f"回测报告已保存: {filename}")


def main():
    """主入口"""
    import argparse

    parser = argparse.ArgumentParser(description='OKX头皮策略交易机器人')
    parser.add_argument('--mode', choices=['backtest', 'live', 'simulate'],
                       default='backtest', help='运行模式')
    parser.add_argument('--config', default=None, help='配置文件路径')

    args = parser.parse_args()

    # 初始化机器人
    bot = ScalpingBot(args.config)

    # 根据模式运行
    if args.mode == 'backtest':
        bot.run_backtest()
    elif args.mode == 'live':
        bot.run_live()
    elif args.mode == 'simulate':
        bot.run_simulate()


if __name__ == '__main__':
    main()
